import os
import sys

import cv2
import numpy as np
import selectivesearch

from tool.MyTool import getRandomInt
from tool.MyTool import remove_salt

sys.path.append('..')


def filter_out(src_frame):
    if src_frame is not None:
        hsv = cv2.cvtColor(src_frame, cv2.COLOR_BGR2HSV)
        lower_val = np.array([0, 0, 0])
        upper_val = np.array([179, 255, 127])
        mask = cv2.inRange(hsv, lower_val, upper_val)
        mask_inv = cv2.bitwise_not(mask)
        return mask_inv


rootPath="./labeled"

files = os.listdir(rootPath)

total = 0
success = 0

for f in files:

    if not f.endswith(".png"):
        continue

    img = cv2.imread(rootPath+os.sep + f)
    total=total+1

    cv2.imshow("source", img)

    cv2.threshold(img, 180, 255, cv2.THRESH_BINARY, img)
    t_width = img.shape[1]
    t_height = img.shape[0]

    cv2.imshow("after threshold", img)

    black_img = filter_out(img.copy())

    cv2.imshow("get black", black_img)

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))

    black_img = cv2.dilate(black_img, kernel)#膨胀

    black_img = cv2.erode(black_img, kernel)#腐蚀
    black_img = cv2.erode(black_img, kernel)

    black_img = remove_salt(black_img, 0, 255, 255, 8)
    black_img = remove_salt(black_img, 0, 255, 255, 8)

    # black_img = cv2.erode(black_img, kernel)

    cv2.imshow("get dilate", black_img)
    cv2.imwrite("temp.png", black_img)
    bgrimg = cv2.imread("temp.png")

    cv2.imshow("reread", bgrimg)

    black_img = bgrimg

    useOpencv = True

    if (useOpencv):
        img_cpy = img.copy()
        img = cv2.cvtColor(black_img, cv2.COLOR_BGR2GRAY)
        contours, hierarchy = cv2.findContours(img, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

        right_counnt = 0

        for con in contours:
            box = cv2.boundingRect(con)  # x，y，w，h
            if box[2] < t_width and (box[3] > t_height * 1.0 / 3):
                # 开始的y坐标:结束的y坐标,开始x:结束的x
                cv2.rectangle(black_img, (box[0], box[1]), (box[0] + box[2], box[1] + box[3]), (0, 0, 255), thickness=1)
                saved = black_img[box[1]:box[1] + box[3], box[0]:box[0] + box[2]]
                cv2.imwrite("./split_data/" + str(getRandomInt()) + ".png", saved)
                right_counnt = right_counnt + 1

        cv2.imshow("x", black_img)
        cv2.waitKey(1000)

        if right_counnt == 4:
            # print("true")
            success = success + 1
        #     cv2.imshow("x", img_cpy)
        #     cv2.waitKey(1000)
        # else:
        #     ""
            # print(f)
            # cv2.imshow("error",img)
            # cv2.waitKey(1000*5)

    else:
        # black_img.shape[0]=3
        # print("shape.channes",black_img.shape[2])

        img_lbl, regions = selectivesearch.selective_search(bgrimg, scale=500, sigma=0.9, min_size=10)

        t_width = img.shape[1]
        t_height = img.shape[0]

        print("width:", str(t_width), " height:" + str(t_height))

        for reg in regions:
            rects = reg['rect']
            width = rects[2]
            height = rects[3]
            print("target width:" + str(width) + " targetHeight:" + str(height))
            if height > t_height / 2 and width < t_width / 4:
                cv2.rectangle(img, (rects[0], rects[1]), (rects[0] + rects[2], rects[1] + rects[3]), (0, 0, 255),
                              thickness=1)
        cv2.imshow("boxed ", img)
        cv2.waitKey(1000)

os.remove("./temp.png")
print(success, "/", total)
